//      graph_main.cc
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp> 
#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/accumulators/statistics/mean.hpp>
#include <boost/foreach.hpp> 
#include <map>

#include "opencv2/imgproc/imgproc.hpp"
#include "opencv2/highgui/highgui.hpp"

#include "stl-helper.h"
#include "data.h"

namespace po = boost::program_options;
namespace boost_acc = boost::accumulators;
po::options_description desc("Allowed options");
po::variables_map vm;

void process_options(int argc, char** argv);
void usage(int exitcode = 0);

void read_data_test();
void icm_test();
void bp_test();

int main(int argc, char** argv)
{
  //MSG("Matching Pursuit for General Loss ... ");

  process_options(argc, argv);

  boost::posix_time::ptime time_start(boost::posix_time::microsec_clock::local_time());

  if (vm.count("help"))
    usage();
  else if (vm.count("read_data"))
    read_data_test();
  else if (vm.count("icm"))
    icm_test();
  else if (vm.count("bp"))
    bp_test();
  else
    usage();

  boost::posix_time::ptime time_end(boost::posix_time::microsec_clock::local_time());
  boost::posix_time::time_duration duration(time_end - time_start);
  MSG("Duration = " << duration);
}

void process_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce this usage")
    ("input", po::value<std::string>(), "input file")
    ("output", po::value<std::string>(), "output file")
    ("verbose", "print debug information")
		("imshow", "show image")

    // OPTIMIZATION PARAMETERS

    // ALGORITHM PARAMETERS
    ("noise", po::value<double>()->default_value(0.1), "noise level")
    // ICM
    ("h", po::value<double>()->default_value(0), "pixel value bias (x)")
    ("beta", po::value<double>()->default_value(1.0), "smoothness (x-x)")
    ("eta", po::value<double>()->default_value(2.0), "correlation (x-y)")
    ("iter", po::value<int>()->default_value(1000), "max #iter")

    //JOB TYPE
    ("read_data", "read data file\n--input --output --noise")
    ("icm", "Iterative conditional mode algorithm\n--input --output ")
    ("bp", "belief propagation\n--input --output ")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}

void usage(int exitcode)
{
  std::cout << desc << std::endl;
  exit(exitcode);
}

void read_data_test()
{
  std::string inFile = VM(vm, "input", std::string);
  std::string outFile = VM(vm, "output", std::string);
  double noise_level = VM(vm, "noise", double);

  cv::Mat gray, fl ;

  gray = cv::imread(inFile, 0); // force GRAY mode
  if(gray.empty()) {
    MSG("Cannot read image file: " << inFile);
    return;
  }
  
  MSG("type=" << gray.type() << " CV_8UC1=" << CV_8UC1);

  arma::vec x = IMG::gray2vec(gray), v = IMG::vec_threshold(x, 0.5, false);
  MSG("x.length=" << sz(x) << " rows=" << gray.rows << " cols=" << gray.cols);

  for (int i = 0; i < sz(v); i++) 
    if (arma::randu() < noise_level) v(i) = -v(i);
  
  IMG::vec2bw(v, fl, gray.size());

  std::vector<int> compression_params;
  compression_params.push_back(CV_IMWRITE_PNG_COMPRESSION);
  compression_params.push_back(9);
  
  try {
    fl *= 255;
    cv::imwrite(outFile, fl, compression_params);
  }
  catch (std::runtime_error& ex) {
    fprintf(stderr, "Exception converting image to PNG format: %s\n", ex.what());
    return;
  }

  cv::namedWindow("Edge map", 1);
  cv::imshow("Edge map", fl);
  cv::waitKey(0);
  
  // cv::Mat_<uchar>::iterator it = gray.begin<uchar>(), itEnd = gray.end<uchar>();
  // for (; it != itEnd; ++it) {
  //   (*it) = 255;
  // }

  // for (int i = 0; i < gray.rows; i++) 
  //   for (int j = 0; j < gray.cols; j++)
  //     gray.at<uchar>(i, j) = 255;

  // cv::namedWindow("Edge map", 1);
  // cv::imshow("Edge map", gray);
  // cv::waitKey(0);
  
}

double local_energy(const vd_t& x, const vd_t& y, int i, double x_i, 
	      int rows, int cols, double h, double beta, double eta)
{
  double e = 0;
  e += h*x_i;
  e += -eta*x_i*y(i);
  int r = i / cols, c = i % cols;
  if (r > 0) e += -beta*x_i*x(i-cols);
  if (c > 0) e += -beta*x_i*x(i-1);
  if (r < rows-1) e += -beta*x_i*x(i+cols);
  if (c > cols-1) e += -beta*x_i*x(i+1);
  return e;
}

void icm_test()
{
  std::string inFile = VM(vm, "input", std::string);
  std::string outFile = VM(vm, "output", std::string);
  double h = VM(vm, "h", double);
  double beta = VM(vm, "beta", double);
  double eta = VM(vm, "eta", double);

  cv::Mat gray, fl ;
  gray = cv::imread(inFile, 0); // force GRAY mode
  if(gray.empty()) {
    MSG("Cannot read image file: " << inFile);
    return;
  }
  cv::imshow("Original", gray);

  arma::vec x = IMG::gray2vec(gray), y = IMG::vec_threshold(x, 0.5, true);
  x = y; // init as noisy image
  int max_iter = 100;
  for (int iter = 0; iter < max_iter; iter++) {
    int count = 0;
    for (int i = 0; i < sz(x); i++) {
      double ePos = local_energy(x, y, i,  1, gray.rows, gray.cols, h, beta, eta);
      double eNeg = local_energy(x, y, i, -1, gray.rows, gray.cols, h, beta, eta);
      double x_i = ePos < eNeg ? 1 : -1;
      if (x_i != x(i)) count++;
      x(i) = x_i;
    }
    MSG("iter=" << iter << " change=" << count);
    if (count == 0) break;
  }

  IMG::vec2bw(x, fl, gray.size());
  cv::namedWindow("Denoise (ICM)", 1);
  cv::imshow("Denoise (ICM)", fl);

  std::vector<int> compression_params;
  compression_params.push_back(CV_IMWRITE_PNG_COMPRESSION);
  compression_params.push_back(9);
  
  try {
    fl *= 255;
    cv::imwrite(outFile, fl, compression_params);
  }
  catch (std::runtime_error& ex) {
    fprintf(stderr, "Exception converting image to PNG format: %s\n", ex.what());
    return;
  }

  cv::waitKey(0);
}

typedef std::pair<int, int> MessageIdx;
typedef std::map<MessageIdx, arma::vec> MessageMap;
typedef MessageMap::value_type MessageType;

class NoiseFactor
{
public:
  int i_, j_;
  int id_;
  static int total_id_;
  std::vector<vi_t>* factor_graph_;
  MessageMap* message_;
  NoiseFactor(int i, int j, int& id, std::vector<vi_t>& fg, MessageMap& message) 
    : i_(i), j_(j), id_(id++), factor_graph_(&fg), message_(&message)
  {
    add_edge();
    add_message(i_, id_);
    add_message(j_, id_);
    //MSG(i_ << " " << j_ << " " << id_);
  }
  void add_edge() {
    factor_graph_->pb(vi_t());
    if (i_ >= 0) {
      (*factor_graph_)[i_].pb(id_);
      (*factor_graph_).back().pb(i_);
    }
    if (j_ >= 0) {
      (*factor_graph_)[j_].pb(id_);
      (*factor_graph_).back().pb(j_);
    }
  }
  void add_message(int i, int f) {
    if (i >= 0) (*message_)[MessageIdx(i,f)] = arma::zeros<arma::vec>(2);
    if (i >= 0) (*message_)[MessageIdx(f,i)] = arma::zeros<arma::vec>(2);
  }  
  static void init_graph(int n, std::vector<vi_t>& fg) {
    fg = std::vector<vi_t>(n);
  }
}; 
int NoiseFactor::total_id_ = 0;

class NoiseBP {
public:
  const arma::vec& y;
  std::vector<vi_t> fg;
  std::vector<NoiseFactor> factor;
  MessageMap message;
  double h, beta, eta;
  int n;

  NoiseBP (const arma::vec& y_)
    : y(y_)
  {
    h = VM(vm, "h", double);
    beta = VM(vm, "beta", double);
    eta = VM(vm, "eta", double);    
    n = sz(y);
  }

  void create_factor_graph(int rows, int cols)
  {
    NoiseFactor::init_graph(n, fg);
    std::vector<NoiseFactor>& f = factor;
    MessageMap& m = message;
    int id = n;
    for (int i = 0; i < n; i++) {
      f.pb(NoiseFactor(i, -1, id, fg, m));

      int r = i / cols, c = i % cols;
      if (r > 0) f.pb(NoiseFactor(i-cols, i, id, fg, m));
      //if (r < rows-1) f.pb(NoiseFactor(i, i+cols, id, fg, m));
      if (c > 0) f.pb(NoiseFactor(i-1, i, id, fg, m));
      //if (c < cols-1) f.pb(NoiseFactor(i, i+1, id, fg, m));
    }
  }

  bool bp(MessageType& m) {
    int i = m.first.first, f = m.first.second;
    arma::vec &msg = m.second;
    arma::vec tmp = arma::zeros<arma::vec>(2);
    arma::vec val(2); val(0) = -1; val(1) = +1;
    if (i < n) { // m(i->f) 
      NoiseFactor ft = factor[f-n]; 
      //MSG("i=" << i << "-->f=" << f << "(" << ft.i_ << "," << ft.j_ << ")");
      if (f != ft.id_) FATAL_ERROR("Check factor id " << f << " != " << ft.id_);
	
      BOOST_FOREACH(int f_j, fg[i]) 
	if (f_j != f) {
	  const arma::vec& m_f_j_i = message[MessageIdx(f_j, i)];
	  if (sz(m_f_j_i) != 2) FATAL_ERROR("Check " << f_j << ", " << i);
	  tmp += m_f_j_i;
	}
    }
    else { // m(f->i)
      std::swap(i, f);
      NoiseFactor ft = factor[f-n];
      //MSG("f=" << f  << "(" << ft.i_ << "," << ft.j_ << ")" << "-->i=" << i);
      if (f != ft.id_) FATAL_ERROR("Check factor id " << f << " != " << ft.id_);

      if (ft.i_ >= 0 && ft.j_ >= 0) {
	int j = ft.j_;
	if (ft.i_ != i) j = ft.i_;

	const arma::vec& m_j_f = message[MessageIdx(j, f)];
	//MSG("m_j_f=" << m_j_f);
	if (sz(m_j_f) != 2) FATAL_ERROR("Check " << i << ", " << j);
	arma::vec neg = -beta*val*(-1) + m_j_f(0);
	arma::vec pos = -beta*val*1 + m_j_f(1);
	tmp(0) = std::min(neg(0), pos(0));
	tmp(1) = std::min(neg(1), pos(1));
      }
      else {
	tmp = h*val - eta*y(i)*val; // f_i(x_i)	 
      }
    }

    bool changed = arma::norm(tmp-msg, 2) > 1e-10;
    msg = tmp;
    MSG(ENDL<<msg);
    return changed;
  }

  void max_sum(arma::vec& x) {
    x = arma::vec(n);
    // max-sum
    for (int i = 0; i < n; i++) {
      arma::vec tmp = arma::zeros<arma::vec>(2);
      BOOST_FOREACH(int f_j, fg[i]) {
	const arma::vec& m_f_j_i = message[MessageIdx(f_j, i)];
	if (sz(m_f_j_i) != 2) FATAL_ERROR("Check " << f_j << ", " << i);
	tmp += m_f_j_i;
      }
      //MSG("i=" << i << " tmp(i)=" << tmp);
			std::cerr << ( (y(i) == 1) ? tmp(0) - tmp(1) : tmp(1) - tmp(0)) << " ";
      x(i) = (tmp(0) < tmp(1)) ? -1 : 1;
    }
		std::cerr << ENDL;
  }

  void add_pending(const MessageIdx& midx, std::set<MessageIdx>& pending)
  {
    int i = midx.first, f = midx.second;
    BOOST_FOREACH(int j, fg[f]) 
      if (j != i) pending.insert(MessageIdx(f, j));

    if (i < n) {
      NoiseFactor ft = factor[f-n];
      if (ft.i_ < 0 || ft.j_ < 0) {
	pending.insert(MessageIdx(f, i));
      }
    }
  }

  double energy(const arma::vec& x) {
    double e = 0;
    BOOST_FOREACH(const NoiseFactor& ft, factor) {
      if (ft.i_ < 0 || ft.j_ < 0)
	e += h*x(ft.i_) - eta*y(ft.i_)*x(ft.i_);
      else
	e += -beta*x(ft.i_)*x(ft.j_);
      std::cerr << e << " ";
    }
    std::cerr << ENDL;
    return e;
  }
};

void bp_test()
{
  std::string inFile = VM(vm, "input", std::string);
  std::string outFile = VM(vm, "output", std::string);
  cv::Mat gray, fl ;
  gray = cv::imread(inFile, 0); // force GRAY mode
  if(gray.empty()) {
    MSG("Cannot read image file: " << inFile);
    return;
  }
  if (vm.count("imshow")) cv::imshow("Original", gray);

  arma::vec x = IMG::gray2vec(gray), y = IMG::vec_threshold(x, 0.5, true);

  NoiseBP nbp(y);
  MSG("Create graph ...");
  nbp.create_factor_graph(gray.rows, gray.cols);
  
  int n = sz(y);
  MSG("n=" << n << " n_factor=" << sz(nbp.factor) << " n_msg=" << sz(nbp.message));

  int max_iter = VM(vm, "iter", int);
  std::set<MessageIdx> pending;
  
  BOOST_FOREACH(const NoiseFactor& ft, nbp.factor) {
    if (ft.i_ < 0 || ft.j_ < 0) {
      MessageIdx midx = MessageIdx(ft.id_, ft.i_);
      MessageType& m = *(nbp.message.find(midx));
      nbp.bp(m);
      nbp.add_pending(midx, pending);
      break;
    }
  }

  for (int iter = 0; iter < max_iter; iter++) {
    int count = 0;
    std::set<MessageIdx> new_pending;
    BOOST_FOREACH(const MessageIdx& midx, pending) {
      MessageType& m = *(nbp.message.find(midx));
      bool changed = nbp.bp(m);
      if (changed) {
	count++;
	nbp.add_pending(midx, new_pending);
      }
    }
    pending = new_pending;
    
    nbp.max_sum(x);
    MSG("iter=" << iter << " change=" << count << " energy="<< nbp.energy(x) << ENDL <<
	"x=" << arma::trans(x) << "y=" << arma::trans(y));

    IMG::vec2bw(x, fl, gray.size());
    
		if (vm.count("imshow")) cv::imshow("Denoise (ICM)", fl);
    if (count == 0) break;
    cv::waitKey(10);
  }

  nbp.max_sum(x);

  IMG::vec2bw(x, fl, gray.size());
  if (vm.count("imshow")) cv::imshow("Denoise (ICM)", fl);

  std::vector<int> compression_params;
  compression_params.push_back(CV_IMWRITE_PNG_COMPRESSION);
  compression_params.push_back(9);
  
  try {
    fl *= 255;
    cv::imwrite(outFile, fl, compression_params);
  }
  catch (std::runtime_error& ex) {
    fprintf(stderr, "Exception converting image to PNG format: %s\n", ex.what());
    return;
  }
  //MSG("y=" << y << "x=" << x);
  cv::waitKey(0);
}
